{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KSIJSS6Wd8t1"
      },
      "source": [
        "##### Copyright 2025 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dOG0dm7LHmBE"
      },
      "source": [
        " <table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Gemma/[Gemma_3n]Finetuned_LoRA_Unsloth_on_Mental_Health_dataset.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sZ-37hWA58vZ"
      },
      "source": [
        "# Gemma-3N Mental Health & Emotional First Aid Assistant for offline inference\n",
        "This notebook fine-tunes Gemma-3N (4B) model using Unsloth on mental health counseling conversations to create an emotional first aid assistant, locally.\n",
        "\n",
        "**Author**: Nguyen Khanh Linh  \n",
        "**GitHub**: [github.com/linhkid](https://github.com/linhkid)  \n",
        "**LinkedIn**: [@Khanh Linh Nguyen](https://twitter.com/https://www.linkedin.com/in/linhnguyenkhanh/)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vCo3nM2358vc"
      },
      "source": [
        "### Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "Hk8M6rW658vd"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "import os, re\n",
        "if \"COLAB_\" not in \"\".join(os.environ.keys()):\n",
        "    !pip install unsloth\n",
        "else:\n",
        "    # Do this only in Colab notebooks! Otherwise use pip install unsloth\n",
        "    import torch; v = re.match(r\"[0-9\\.]{3,}\", str(torch.__version__)).group(0)\n",
        "    xformers = \"xformers==\" + (\"0.0.32.post2\" if v == \"2.8.0\" else \"0.0.29.post3\")\n",
        "    !pip install --no-deps bitsandbytes accelerate {xformers} peft trl triton cut_cross_entropy unsloth_zoo\n",
        "    !pip install sentencepiece protobuf \"datasets>=3.4.1,<4.0.0\" \"huggingface_hub>=0.34.0\" hf_transfer\n",
        "    !pip install --no-deps unsloth\n",
        "!pip install transformers==4.55.4\n",
        "!pip install --no-deps trl==0.22.2\n",
        "import torch; torch._dynamo.config.recompile_limit = 64;"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "zQ7ewecW58vg"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install --no-deps --upgrade timm # Only for Gemma 3N"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Ny7mplo58vh"
      },
      "source": [
        "### Load Gemma-3N Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "xd0cI7Py58vi"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
            "🦥 Unsloth Zoo will now patch everything to make training faster!\n",
            "==((====))==  Unsloth 2025.10.1: Fast Gemma3N patching. Transformers: 4.55.4.\n",
            "   \\\\   /|    NVIDIA L4. Num GPUs = 1. Max memory: 22.161 GB. Platform: Linux.\n",
            "O^O/ \\_/ \\    Torch: 2.8.0+cu126. CUDA: 8.9. CUDA Toolkit: 12.6. Triton: 3.4.0\n",
            "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.32.post2. FA2 = False]\n",
            " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
            "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "8e2744fdbb624ad1b5a5122c7e7e3be2",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model.safetensors.index.json: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "c4e98f4d1c4743f3b555766bba6e8d0c",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Fetching 3 files:   0%|          | 0/3 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "ba59de2ba16e4fd0ab2c5f26d3516515",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model-00003-of-00003.safetensors:   0%|          | 0.00/1.15G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "e4e6e18b8bd341bcba2e42ef838abda5",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model-00001-of-00003.safetensors:   0%|          | 0.00/3.72G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "08f37d51428c47f589ce5d351b10cdaa",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model-00002-of-00003.safetensors:   0%|          | 0.00/4.99G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "8e1908eeb419455ca71ce335e5df7fc0",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "c6e595efb420480cb1445e657bf15326",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "generation_config.json:   0%|          | 0.00/210 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "63e06f41b2d54a50961e287b8578cb8c",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "processor_config.json:   0%|          | 0.00/98.0 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "19142172307b4f6887271b287d7200cf",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "chat_template.jinja: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "bde55ba01fe04bcf80c0ea9840aaf22d",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "preprocessor_config.json: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "f215eec60cc04f56b2d7442f87f36044",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "tokenizer_config.json: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "34b232557aea4cd39add7779f66bab90",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "tokenizer.model:   0%|          | 0.00/4.70M [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "6772cd85a4774986ad212c2aa58e8362",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "tokenizer.json:   0%|          | 0.00/33.4M [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "56934dd41e9c432dbbb621fc1b7d3216",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "special_tokens_map.json:   0%|          | 0.00/777 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from unsloth import FastModel\n",
        "import torch\n",
        "\n",
        "model, tokenizer = FastModel.from_pretrained(\n",
        "    model_name = \"unsloth/gemma-3n-E4B-it\",\n",
        "    dtype = None, # None for auto detection\n",
        "    max_seq_length = 2048, # Increased for longer conversations\n",
        "    load_in_4bit = True,  # 4 bit quantization to reduce memory\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_mZd14aN58vj"
      },
      "source": [
        "### Load Mental Health Datasets\n",
        "We'll combine two datasets:\n",
        "1. Mental health counseling conversations\n",
        "2. Psychology Q&A dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "sHVb8Uil58vl"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading datasets...\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "8f808b54f0fc48f987fd39b08d07d623",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "README.md: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9c3463b72dcb402396f52394e3b796fa",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "combined_dataset.json: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "284bfed2607242b1b55f2f53a40a5df2",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Generating train split:   0%|          | 0/3512 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "ee2d70ecfbd34c54a9f29aa43d0f73ec",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "data/part-00000-694db9fd-774c-4205-b938-(…):   0%|          | 0.00/1.59M [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "b460efa9b05540b4a9282f8fe548a6cc",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "data/part-00001-694db9fd-774c-4205-b938-(…):   0%|          | 0.00/96.4k [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "dfd71f383285478599aebba2cb4e41bf",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Generating train split:   0%|          | 0/9846 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Datasets loaded.\n",
            "Combined dataset has 13358 samples.\n",
            "Training dataset has 12022 samples.\n",
            "Evaluation dataset has 1336 samples.\n",
            "\n",
            "Sample of raw data:\n",
            "Context: I'm having trouble with my eating habits. I feel like I'm always bingeing and then feeling guilty about it....\n",
            "Response: Eating disorders can be serious, but there is help available. Let's work on developing healthy eating habits, exploring possible underlying causes for your behavior, and possibly seeking out therapy o...\n"
          ]
        }
      ],
      "source": [
        "from datasets import load_dataset, concatenate_datasets\n",
        "\n",
        "print(\"Loading datasets...\")\n",
        "\n",
        "# Load the first dataset\n",
        "mental_health_dataset = load_dataset(\"Amod/mental_health_counseling_conversations\", split=\"train\")\n",
        "\n",
        "# Load the second dataset and rename columns to match the first\n",
        "psychology_dataset = load_dataset(\"jkhedri/psychology-dataset\", split=\"train\")\n",
        "psychology_dataset = psychology_dataset.rename_columns({\n",
        "    \"question\": \"Context\",\n",
        "    \"response_j\": \"Response\"\n",
        "})\n",
        "\n",
        "# Combine the datasets\n",
        "combined_dataset = concatenate_datasets([\n",
        "    mental_health_dataset.select_columns(['Context', 'Response']),\n",
        "    psychology_dataset.select_columns(['Context', 'Response'])\n",
        "])\n",
        "\n",
        "# Split the combined dataset into train and evaluation sets\n",
        "dataset = combined_dataset.train_test_split(test_size=0.1, seed=42) # Using 10% for evaluation\n",
        "\n",
        "print(\"Datasets loaded.\")\n",
        "print(f\"Combined dataset has {len(combined_dataset)} samples.\")\n",
        "print(f\"Training dataset has {len(dataset['train'])} samples.\")\n",
        "print(f\"Evaluation dataset has {len(dataset['test'])} samples.\")\n",
        "\n",
        "print(\"\\nSample of raw data:\")\n",
        "print(f\"Context: {dataset['train'][0]['Context'][:200]}...\")\n",
        "print(f\"Response: {dataset['train'][0]['Response'][:200]}...\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eVCP8oGx58vm"
      },
      "source": [
        "### Setup Chat Template for Gemma-3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "rF7qjxqD58vn"
      },
      "outputs": [],
      "source": [
        "from unsloth.chat_templates import get_chat_template\n",
        "tokenizer = get_chat_template(\n",
        "    tokenizer,\n",
        "    chat_template = \"gemma-3\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HZo6UGtO58vn"
      },
      "source": [
        "### Format Dataset with Chat Template\n",
        "Convert Context/Response pairs into Gemma-3 conversation format"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "WOCTmTI558vo"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a350a4a2273449b68826d1ec2bf543c7",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map:   0%|          | 0/12022 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1c7db332adbb4dde8742780825e9e2ec",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map:   0%|          | 0/1336 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Datasets loaded and formatted.\n",
            "Combined dataset has 2 samples.\n",
            "\n",
            "Sample of formatted data:\n",
            "<start_of_turn>user\n",
            "I'm having trouble with my eating habits. I feel like I'm always bingeing and then feeling guilty about it.<end_of_turn>\n",
            "<start_of_turn>model\n",
            "Eating disorders can be serious, but there is help available. Let's work on developing healthy eating habits, exploring possible underlying causes for your behavior, and possibly seeking out therapy or support groups.<end_of_turn>\n",
            "\n"
          ]
        }
      ],
      "source": [
        "def formatting_prompts_func(examples):\n",
        "    convos = []\n",
        "    for context, response in zip(examples[\"Context\"], examples[\"Response\"]):\n",
        "        messages = [\n",
        "            {\"role\": \"user\", \"content\": context},\n",
        "            {\"role\": \"assistant\", \"content\": response},\n",
        "        ]\n",
        "        # add_generation_prompt=False to avoid a user turn at the end\n",
        "        # Remove <bos> token as processor will add it\n",
        "        text = tokenizer.apply_chat_template(\n",
        "            messages,\n",
        "            tokenize=False,\n",
        "            add_generation_prompt=False\n",
        "        ).removeprefix('<bos>')\n",
        "        convos.append(text)\n",
        "    return {\"text\": convos}\n",
        "\n",
        "# Apply the formatting function to both train and test splits\n",
        "dataset[\"train\"] = dataset[\"train\"].map(formatting_prompts_func, batched=True)\n",
        "dataset[\"test\"] = dataset[\"test\"].map(formatting_prompts_func, batched=True)\n",
        "\n",
        "print(\"\\nDatasets loaded and formatted.\")\n",
        "print(f\"Combined dataset has {len(dataset)} samples.\")\n",
        "# Optional: Print a sample to verify the format\n",
        "print(\"\\nSample of formatted data:\")\n",
        "print(dataset[\"train\"][0][\"text\"][:500])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "90fmDsac58vo"
      },
      "source": [
        "### Add LoRA Adapters\n",
        "We only update a small amount of parameters for efficient fine-tuning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "SJ_gCbO_58vp"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Unsloth: Making `model.base_model.model.model.language_model` require gradients\n"
          ]
        }
      ],
      "source": [
        "model = FastModel.get_peft_model(\n",
        "    model,\n",
        "    finetune_vision_layers     = False, # Turn off for just text!\n",
        "    finetune_language_layers   = True,  # Should leave on!\n",
        "    finetune_attention_modules = True,  # Attention good for GRPO\n",
        "    finetune_mlp_modules       = True,  # Should leave on always!\n",
        "\n",
        "    r = 16,          # Larger = higher accuracy, increased for mental health\n",
        "    lora_alpha = 16, # Recommended alpha == r at least\n",
        "    lora_dropout = 0,\n",
        "    bias = \"none\",\n",
        "    random_state = 3407,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BJVj7iBY58vq"
      },
      "source": [
        "### Setup Training Configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "Y7i_eZ2958vq"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "aeaf6ba0f5804b7daf3cfb2d82acb9c1",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Unsloth: Tokenizing [\"text\"] (num_proc=16):   0%|          | 0/12022 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "ba091806ec0946f9a36cc22bfb79b7c7",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Unsloth: Tokenizing [\"text\"] (num_proc=16):   0%|          | 0/1336 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from trl import SFTTrainer, SFTConfig\n",
        "\n",
        "trainer = SFTTrainer(\n",
        "    model = model,\n",
        "    tokenizer = tokenizer,\n",
        "    train_dataset = dataset[\"train\"],\n",
        "    eval_dataset = dataset[\"test\"], # Can set up evaluation!\n",
        "    args = SFTConfig(\n",
        "        dataset_text_field = \"text\",\n",
        "        per_device_train_batch_size = 4,\n",
        "        gradient_accumulation_steps = 4,\n",
        "        warmup_steps = 10,\n",
        "        #num_train_epochs = 1, # Full training run\n",
        "        max_steps = 100, # Or use max_steps for quick testing\n",
        "        learning_rate = 2e-4,\n",
        "        logging_steps = 50,\n",
        "        optim = \"adamw_8bit\",\n",
        "        weight_decay = 0.01,\n",
        "        lr_scheduler_type = \"linear\",\n",
        "        seed = 3407,\n",
        "        output_dir = \"outputs\",\n",
        "        report_to = \"none\", # Use this for WandB etc\n",
        "        gradient_checkpointing = True, # Disable gradient checkpointing as a workaround\n",
        "    ),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VkAbegxy58vq"
      },
      "source": [
        "### Train Only on Assistant Responses\n",
        "This helps increase accuracy by masking out the user inputs during training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "FHDL94Nq58vr"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9f152a58ef934bb3a17efb76fd870783",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map (num_proc=12):   0%|          | 0/12022 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "24671ee03a664629919b16545d244490",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map (num_proc=12):   0%|          | 0/1336 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from unsloth.chat_templates import train_on_responses_only\n",
        "\n",
        "trainer = train_on_responses_only(\n",
        "    trainer,\n",
        "    instruction_part = \"<start_of_turn>user\\n\",\n",
        "    response_part = \"<start_of_turn>model\\n\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JdeK2KEd58vr"
      },
      "source": [
        "### Verify Masking\n",
        "Let's check that the instruction masking is working correctly"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "hu_8QlHM58vr"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Full conversation:\n",
            "<bos><start_of_turn>user\n",
            "I'm struggling with grief and loss.<end_of_turn>\n",
            "<start_of_turn>model\n",
            "Grief and loss can be a difficult experience, but it's possible to work through the emotions and find healing. Let's work together to identify any underlying beliefs or experiences that may be contributing to your grief and develop strategies to manage the emotions. It may also be helpful to seek out support from loved ones or grief-specific therapy or support groups.<end_of_turn>\n",
            "\n",
            "\n",
            "==================================================\n",
            "\n",
            "Masked (only assistant response will be trained):\n",
            "                  Grief and loss can be a difficult experience, but it's possible to work through the emotions and find healing. Let's work together to identify any underlying beliefs or experiences that may be contributing to your grief and develop strategies to manage the emotions. It may also be helpful to seek out support from loved ones or grief-specific therapy or support groups.<end_of_turn>\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(\"Full conversation:\")\n",
        "print(tokenizer.decode(trainer.train_dataset[5][\"input_ids\"]))\n",
        "print(\"\\n\" + \"=\"*50 + \"\\n\")\n",
        "print(\"Masked (only assistant response will be trained):\")\n",
        "print(tokenizer.decode([tokenizer.pad_token_id if x == -100 else x for x in trainer.train_dataset[5][\"labels\"]]).replace(tokenizer.pad_token, \" \"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tUZJ_-5358vs"
      },
      "source": [
        "### Check Memory Stats"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "vS_tCpGp58vs"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GPU = NVIDIA L4. Max memory = 22.161 GB.\n",
            "9.379 GB of memory reserved.\n"
          ]
        }
      ],
      "source": [
        "gpu_stats = torch.cuda.get_device_properties(0)\n",
        "start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
        "max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)\n",
        "print(f\"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.\")\n",
        "print(f\"{start_gpu_memory} GB of memory reserved.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UX_eACi958vt"
      },
      "source": [
        "### Train the Model!\n",
        "To resume a training run, set `trainer.train(resume_from_checkpoint = True)`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "Mi1watL758vt"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1\n",
            "   \\\\   /|    Num examples = 12,022 | Num Epochs = 1 | Total steps = 100\n",
            "O^O/ \\_/ \\    Batch size per device = 4 | Gradient accumulation steps = 4\n",
            "\\        /    Data Parallel GPUs = 1 | Total batch size (4 x 4 x 1) = 16\n",
            " \"-____-\"     Trainable parameters = 38,420,480 of 7,888,398,672 (0.49% trained)\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Unsloth: Will smartly offload gradients to save VRAM!\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='100' max='100' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [100/100 09:26, Epoch 0/1]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>50</td>\n",
              "      <td>2.631100</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>100</td>\n",
              "      <td>1.984400</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "trainer_stats = trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5n2p9Eo-58vt"
      },
      "source": [
        "### Training Stats"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "Wsr2P9Q_58vt"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "652.2686 seconds used for training.\n",
            "10.87 minutes used for training.\n",
            "Peak reserved memory = 14.281 GB.\n",
            "Peak reserved memory for training = 4.902 GB.\n",
            "Peak reserved memory % of max memory = 64.442 %.\n",
            "Peak reserved memory for training % of max memory = 22.12 %.\n"
          ]
        }
      ],
      "source": [
        "used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
        "used_memory_for_lora = round(used_memory - start_gpu_memory, 3)\n",
        "used_percentage = round(used_memory / max_memory * 100, 3)\n",
        "lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)\n",
        "print(f\"{trainer_stats.metrics['train_runtime']} seconds used for training.\")\n",
        "print(f\"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.\")\n",
        "print(f\"Peak reserved memory = {used_memory} GB.\")\n",
        "print(f\"Peak reserved memory for training = {used_memory_for_lora} GB.\")\n",
        "print(f\"Peak reserved memory % of max memory = {used_percentage} %.\")\n",
        "print(f\"Peak reserved memory for training % of max memory = {lora_percentage} %.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jK4CyMGD58vt"
      },
      "source": [
        "### Test the Mental Health Assistant\n",
        "Let's test the fine-tuned model on mental health queries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "qfUtsU4P58vt"
      },
      "outputs": [],
      "source": [
        "from transformers import TextStreamer\n",
        "\n",
        "def chat_with_mental_health_assistant(user_message, max_new_tokens=256):\n",
        "    messages = [{\n",
        "        \"role\": \"user\",\n",
        "        \"content\": [{\"type\": \"text\", \"text\": user_message}]\n",
        "    }]\n",
        "\n",
        "    inputs = tokenizer.apply_chat_template(\n",
        "        messages,\n",
        "        add_generation_prompt = True,\n",
        "        return_tensors = \"pt\",\n",
        "        tokenize = True,\n",
        "        return_dict = True,\n",
        "    ).to(\"cuda\")\n",
        "\n",
        "    _ = model.generate(\n",
        "        **inputs,\n",
        "        max_new_tokens = max_new_tokens,\n",
        "        temperature = 0.7,  # Lower temperature for more focused responses\n",
        "        top_p = 0.9,\n",
        "        top_k = 50,\n",
        "        streamer = TextStreamer(tokenizer, skip_prompt = True),\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "RiqPA9MF58vu"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "==================================================\n",
            "User: I've been feeling very anxious lately and I don't know why.\n",
            "==================================================\n",
            "It's common to experience anxiety without a clear trigger. Let's explore your recent experiences and see if we can identify any patterns or potential stressors. We can also discuss relaxation techniques and coping strategies to help manage your anxiety.<end_of_turn>\n"
          ]
        }
      ],
      "source": [
        "# Test 1: Anxiety\n",
        "print(\"=\" * 50)\n",
        "print(\"User: I've been feeling very anxious lately and I don't know why.\")\n",
        "print(\"=\" * 50)\n",
        "chat_with_mental_health_assistant(\"I've been feeling very anxious lately and I don't know why.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "_v2OkfQL58vu"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "==================================================\n",
            "User: I'm feeling overwhelmed with work and personal life. What should I do?\n",
            "==================================================\n",
            "It's common to feel overwhelmed with work and personal life at times. Let's work together to identify your priorities and develop a plan to manage your time and energy effectively. This may include delegating tasks, setting boundaries, and practicing self-care.<end_of_turn>\n"
          ]
        }
      ],
      "source": [
        "# Test 2: Stress\n",
        "print(\"\\n\" + \"=\" * 50)\n",
        "print(\"User: I'm feeling overwhelmed with work and personal life. What should I do?\")\n",
        "print(\"=\" * 50)\n",
        "chat_with_mental_health_assistant(\"I'm feeling overwhelmed with work and personal life. What should I do?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "htp0oZiS58vu"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "==================================================\n",
            "User: I've lost interest in things I used to enjoy. Is this normal?\n",
            "==================================================\n",
            "It's normal to experience a temporary loss of interest in things you used to enjoy, especially after a stressful period. However, if this persists for more than a few weeks or is accompanied by other symptoms like fatigue, sleep disturbances, or changes in appetite, it could be a sign of depression. It's important to consult with a healthcare professional for proper diagnosis and treatment.<end_of_turn>\n"
          ]
        }
      ],
      "source": [
        "# Test 3: Depression symptoms\n",
        "print(\"\\n\" + \"=\" * 50)\n",
        "print(\"User: I've lost interest in things I used to enjoy. Is this normal?\")\n",
        "print(\"=\" * 50)\n",
        "chat_with_mental_health_assistant(\"I've lost interest in things I used to enjoy. Is this normal?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "LAEmEngu58vu"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "==================================================\n",
            "User: What are some healthy coping mechanisms for stress?\n",
            "==================================================\n",
            "Healthy coping mechanisms for stress include exercise, spending time in nature, practicing mindfulness and meditation, engaging in hobbies, and connecting with loved ones. It's also important to prioritize sleep and a balanced diet.<end_of_turn>\n"
          ]
        }
      ],
      "source": [
        "# Test 4: Self-care\n",
        "print(\"\\n\" + \"=\" * 50)\n",
        "print(\"User: What are some healthy coping mechanisms for stress?\")\n",
        "print(\"=\" * 50)\n",
        "chat_with_mental_health_assistant(\"What are some healthy coping mechanisms for stress?\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cihpuwme58vv"
      },
      "source": [
        "### Save the Model\n",
        "Save the LoRA adapters locally"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "zb0e1faR58vv"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model saved successfully!\n"
          ]
        }
      ],
      "source": [
        "model.save_pretrained(\"gemma-3n-mental-health\")  # Local saving\n",
        "tokenizer.save_pretrained(\"gemma-3n-mental-health\")\n",
        "print(\"Model saved successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dIwEdO_W58vw"
      },
      "source": [
        "### Optional: Push to Hugging Face Hub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "rMqes-W458vw"
      },
      "outputs": [],
      "source": [
        "# Uncomment and add your token to upload\n",
        "# model.push_to_hub(\"YOUR_HF_USERNAME/gemma-3n-mental-health\", token = \"hf_...\")\n",
        "# tokenizer.push_to_hub(\"YOUR_HF_USERNAME/gemma-3n-mental-health\", token = \"hf_...\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8l0khHJ58vw"
      },
      "source": [
        "### Load Model for Future Use"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "Q5IfDoH858vw"
      },
      "outputs": [],
      "source": [
        "# To load the fine-tuned model later:\n",
        "if False:  # Set to True to load\n",
        "    from unsloth import FastModel\n",
        "    model, tokenizer = FastModel.from_pretrained(\n",
        "        model_name = \"gemma-3n-mental-health\",\n",
        "        max_seq_length = 2048,\n",
        "        load_in_4bit = True,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F36DWsOg58vw"
      },
      "source": [
        "### Export to GGUF (for llama.cpp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "LE8fw4Xa58v4"
      },
      "outputs": [],
      "source": [
        "# Save to GGUF format for local deployment\n",
        "if False:  # Set to True to export\n",
        "    model.save_pretrained_gguf(\n",
        "        \"gemma-3n-mental-health-gguf\",\n",
        "        quantization_type = \"Q8_0\",  # Q8_0, BF16, or F16\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GfKaBYM558v4"
      },
      "source": [
        "### Export to float16 (for deployment)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "lmvamQAS58v4"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "e4b056d097e74a78a22e96ce28b6f9f4",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "config.json: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Found HuggingFace hub cache directory: /root/.cache/huggingface/hub\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9ba74ee0a374461a8dd1b1b58eaf9e2d",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Fetching 1 files:   0%|          | 0/1 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "69675b2ae03940ea91bca29eeb88755c",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model.safetensors.index.json: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Checking cache directory for required files...\n",
            "Cache check failed: model-00001-of-00004.safetensors not found in local cache.\n",
            "Not all required files found in cache. Will proceed with downloading.\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\rUnsloth: Preparing safetensor model files:   0%|          | 0/4 [00:00<?, ?it/s]"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "0e902f15be6c43ae869c81c14235d279",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model-00001-of-00004.safetensors:   0%|          | 0.00/3.08G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\rUnsloth: Preparing safetensor model files:  25%|██▌       | 1/4 [00:10<00:31, 10.66s/it]"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "68cafdb7958543acbdb49e60661ae342",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model-00002-of-00004.safetensors:   0%|          | 0.00/4.97G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\rUnsloth: Preparing safetensor model files:  50%|█████     | 2/4 [00:36<00:39, 19.57s/it]"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "e924716008254ec5946e921f5a2a4321",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model-00003-of-00004.safetensors:   0%|          | 0.00/4.99G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\rUnsloth: Preparing safetensor model files:  75%|███████▌  | 3/4 [02:12<00:54, 54.33s/it]"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "07365c21e7e1446286a4781b1394c4fc",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "model-00004-of-00004.safetensors:   0%|          | 0.00/2.66G [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Unsloth: Preparing safetensor model files: 100%|██████████| 4/4 [02:21<00:00, 35.31s/it]\n",
            "Unsloth: Merging weights into 16bit: 100%|██████████| 4/4 [00:53<00:00, 13.27s/it]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Unsloth: Merge process complete. Saved to `/content/gemma-3n-mental-health-f16`\n"
          ]
        }
      ],
      "source": [
        "# Save merged model in float16\n",
        "if True:  # Set to True to export\n",
        "    model.save_pretrained_merged(\"gemma-3n-mental-health-f16\", tokenizer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J9PDxTbhMFyw"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "torch.cuda.empty_cache()\n",
        "print(\"CUDA cache cleared.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tEUQz9qSEd9U"
      },
      "source": [
        "## Local Inference for fully merged model\n",
        "\n",
        "Only run if you have enough VRAM. Otherwise, you can choose to deploy to cloud"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "d3uinSKKEeO4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
            "🦥 Unsloth Zoo will now patch everything to make training faster!\n",
            "==((====))==  Unsloth 2025.10.1: Fast Gemma3N patching. Transformers: 4.55.4.\n",
            "   \\\\   /|    NVIDIA L4. Num GPUs = 1. Max memory: 22.161 GB. Platform: Linux.\n",
            "O^O/ \\_/ \\    Torch: 2.8.0+cu126. CUDA: 8.9. CUDA Toolkit: 12.6. Triton: 3.4.0\n",
            "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.32.post2. FA2 = False]\n",
            " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
            "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n",
            "Unsloth: QLoRA and full finetuning all not selected. Switching to 16bit LoRA.\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "97cec90dc0fb4402b01694879dae93b3",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "User: I've been feeling very anxious lately and I don't know why.\n"
          ]
        }
      ],
      "source": [
        "from unsloth import FastModel\n",
        "from transformers import TextStreamer\n",
        "import torch # Import torch for dtype\n",
        "\n",
        "# Load the merged model with bfloat16 precision\n",
        "model, tokenizer = FastModel.from_pretrained(\n",
        "    model_name = \"gemma-3n-mental-health-f16\", # Path to the merged model\n",
        "    max_seq_length = 2048,\n",
        "    dtype = torch.bfloat16, # Use bfloat16 for the merged model\n",
        "    load_in_4bit = False, # Not needed for merged model\n",
        ")\n",
        "\n",
        "def chat_with_merged_model(user_message, max_new_tokens=256):\n",
        "    messages = [{\n",
        "        \"role\": \"user\",\n",
        "        \"content\": [{\"type\": \"text\", \"text\": user_message}]\n",
        "    }]\n",
        "\n",
        "    inputs = tokenizer.apply_chat_template(\n",
        "        messages,\n",
        "        add_generation_prompt = True,\n",
        "        return_tensors = \"pt\",\n",
        "        tokenize = True,\n",
        "        return_dict = True,\n",
        "    ).to(\"cuda\") # Move to CUDA for inference\n",
        "\n",
        "    _ = model.generate(\n",
        "        **inputs,\n",
        "        max_new_tokens = max_new_tokens,\n",
        "        temperature = 0.7,\n",
        "        top_p = 0.9,\n",
        "        top_k = 50,\n",
        "        streamer = TextStreamer(tokenizer, skip_prompt = True),\n",
        "    )\n",
        "\n",
        "# Example usage\n",
        "q = \"I've been feeling very anxious lately and I don't know why.\"\n",
        "print(\"User: {}\".format(q))\n",
        "#chat_with_merged_model(q) #uncomment if you can run locally"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "r7Tr-Oflat69"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "It's common to experience anxiety without a clear trigger. Let's explore your life and see if there are any potential stressors or patterns that might be contributing to your anxiety. We can also discuss relaxation techniques and coping strategies to help manage your symptoms.<end_of_turn>\n"
          ]
        }
      ],
      "source": [
        "chat_with_merged_model(q) #uncomment if you can run locally"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "699f18b9"
      },
      "source": [
        "#Custom (does not work rn): Import Model to Vertex AI\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m5enzi1QNku9"
      },
      "source": [
        "## Deploy to Vertex AI as Custom Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "J1T6iD9hNm97"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "95ca7726"
      },
      "source": [
        "## Define GCP Project and Bucket\n",
        "\n",
        "### Subtask:\n",
        "Specify your Google Cloud project ID and a Cloud Storage bucket name where the model will be stored."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "66bd4f37"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"project-id\" # @param {type:\"string\"}\n",
        "BUCKET_NAME = \"bucket-name\" # @param {type:\"string\"}\n",
        "REGION = \"region\" # @param {type:\"string\"}\n",
        "\n",
        "# Ensure the bucket name is valid\n",
        "if not BUCKET_NAME or \"your-gcp-bucket-name\" in BUCKET_NAME:\n",
        "    raise ValueError(\"Please replace 'your-gcp-bucket-name' with your actual bucket name.\")\n",
        "\n",
        "# Create the bucket if it doesn't exist, explicitly setting the project\n",
        "!gsutil -p {PROJECT_ID} mb -l {REGION} gs://{BUCKET_NAME}\n",
        "\n",
        "print(f\"Project ID: {PROJECT_ID}\")\n",
        "print(f\"Bucket Name: {BUCKET_NAME}\")\n",
        "print(f\"Region: {REGION}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9a64a878"
      },
      "source": [
        "## Upload Merged Model to Cloud Storage\n",
        "\n",
        "Upload the saved merged model files (from the `gemma-3n-mental-health-f16` directory) to the specified Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "91705c13"
      },
      "outputs": [],
      "source": [
        "# Upload the merged model directory to the bucket\n",
        "!gsutil -m cp -r gemma-3n-mental-health-f16 gs://{BUCKET_NAME}/gemma-3n-mental-health-f16\n",
        "\n",
        "print(f\"Model uploaded to gs://{BUCKET_NAME}/gemma-3n-mental-health-f16\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "731398db"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "# Initialize Vertex AI\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
        "\n",
        "# Define model parameters\n",
        "MODEL_DISPLAY_NAME = \"gemma-3n-mental-health-assistant\"\n",
        "ARTIFACT_URI = f\"gs://{BUCKET_NAME}/gemma-3n-mental-health-f16\"\n",
        "SERVING_CONTAINER_IMAGE = \"us-docker.pkg.dev/vertex-ai/vertex-ai-huggingface-inference/transformers-tgi-nvidia-l4:latest\" # Example image for L4 GPU\n",
        "\n",
        "# Import the model\n",
        "model = aiplatform.Model.upload(\n",
        "    display_name=MODEL_DISPLAY_NAME,\n",
        "    artifact_uri=ARTIFACT_URI,\n",
        "    serving_container_image_uri=SERVING_CONTAINER_IMAGE,\n",
        "    sync=True, # Wait for the import to complete\n",
        ")\n",
        "\n",
        "print(f\"Model uploaded to Vertex AI: {model.resource_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "17b30ea2"
      },
      "source": [
        "## Create Vertex AI Endpoint\n",
        "\n",
        "Create a Vertex AI Endpoint resource where the model will be deployed for online predictions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6cc45260"
      },
      "outputs": [],
      "source": [
        "# Create an endpoint\n",
        "endpoint = aiplatform.Endpoint.create(\n",
        "    display_name=f\"{MODEL_DISPLAY_NAME}_endpoint\",\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        ")\n",
        "\n",
        "print(f\"Endpoint created: {endpoint.resource_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "694fb39c"
      },
      "source": [
        "## Deploy Model to Endpoint\n",
        "\n",
        "Deploy the imported Vertex AI Model to the created Endpoint. This step involves configuring the deployment, such as specifying the machine type and the number of replicas."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e4f4874d"
      },
      "outputs": [],
      "source": [
        "# Define deployment parameters\n",
        "MACHINE_TYPE = \"g2-standard-8\" # @param {type:\"string\"} # Machine type with L4 GPU\n",
        "ACCELERATOR_TYPE = \"NVIDIA_L4\" # @param {type:\"string\"} # L4 GPU\n",
        "ACCELERATOR_COUNT = 1 # @param {type:\"integer\"}\n",
        "\n",
        "# Deploy the model to the endpoint\n",
        "endpoint.deploy(\n",
        "    model=model,\n",
        "    deployed_model_display_name=MODEL_DISPLAY_NAME,\n",
        "    machine_type=MACHINE_TYPE,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    accelerator_count=ACCELERATOR_COUNT,\n",
        "    sync=True, # Wait for the deployment to complete\n",
        ")\n",
        "\n",
        "print(f\"Model deployed to endpoint: {endpoint.resource_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-K50MZk258v6"
      },
      "source": [
        "## Important Disclaimer\n",
        "\n",
        "⚠️ **This model is for educational and research purposes only.** It should NOT replace professional mental health care.\n",
        "\n",
        "- This AI assistant provides general emotional support and information\n",
        "- It is NOT a substitute for professional therapy or medical advice\n",
        "- In crisis situations, please contact:\n",
        "  - **National Suicide Prevention Lifeline**: 988 (US)\n",
        "  - **Crisis Text Line**: Text HOME to 741741\n",
        "  - **International Association for Suicide Prevention**: https://www.iasp.info/resources/Crisis_Centres/\n",
        "\n",
        "Always consult with qualified mental health professionals for diagnosis and treatment."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "[Gemma_3n]Finetuned_LoRA_Unsloth_on_Mental_Health_dataset.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
